#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
通用数据导入脚本
支持.xls、.xlsx和.csv格式文件的导入
可通过--teacher_col和--title_col参数手动指定列索引
支持通过-o参数指定中间CSV文件输出路径
"""

import pandas as pd
import json
import os
import argparse
from datetime import datetime
import numpy as np
import sys
import faiss
import jieba
import jieba.analyse

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# 自定义同义词词典
SYNONYM_DICT = {
    "深度学习": "DL",
    "深度神经网络": "DNN",
    "卷积神经网络": "CNN",
    "循环神经网络": "RNN",
    "长短期记忆网络": "LSTM",
    "长短时记忆网络": "LSTM",
    "门控循环单元": "GRU",
    "支持向量机": "SVM",
    "随机森林": "RF",
    "自然语言处理": "NLP",
    "计算机视觉": "CV",
    "强化学习": "RL",
    "机器学习": "ML",
    "人工神经网络": "ANN",
    "生成对抗网络": "GAN",
    "变分自编码器": "VAE"
}

# 模型标志
model = None
model_initialized = False

# 论文题目生成相关的数据
TEACHER_NAMES = [
    "张教授", "李教授", "王教授", "赵教授", "陈教授", "刘教授", "杨教授", "黄教授", "周教授", "吴教授",
    "徐教授", "孙教授", "胡教授", "朱教授", "高教授", "林教授", "何教授", "郭教授", "马教授", "罗教授",
    "梁教授", "宋教授", "郑教授", "谢教授", "韩教授", "唐教授", "冯教授", "于教授", "董教授", "萧教授"
]

RESEARCH_FIELDS = [
    "深度学习", "机器学习", "数据挖掘", "自然语言处理", "计算机视觉", "图像识别", "语音识别", 
    "推荐系统", "知识图谱", "数据可视化", "大数据分析", "云计算", "边缘计算", "物联网", "区块链",
    "人工智能", "智能算法", "模式识别", "智能控制", "机器人技术", "虚拟现实", "增强现实", 
    "生物信息学", "医疗图像处理", "金融科技", "量化交易", "网络安全", "信息检索", "数据库技术"
]

RESEARCH_OBJECTS = [
    "图像识别技术", "文本分类方法", "用户行为分析", "推荐算法优化", "情感分析模型", "目标检测技术",
    "语音合成方法", "机器翻译系统", "数据聚类算法", "异常检测技术", "预测模型构建", "特征选择方法",
    "网络爬虫技术", "数据清洗方法", "可视化设计", "分布式计算", "负载均衡策略", "容错机制",
    "加密算法", "身份认证技术", "入侵检测系统", "防火墙技术", "数据压缩方法", "存储优化技术"
]

RESEARCH_METHODS = [
    "研究", "分析", "设计", "实现", "应用", "探讨", "优化", "改进", "开发", "构建", 
    "验证", "评估", "比较", "调查", "实验", "建模", "仿真", "测试"
]

def load_model():
    """加载模型"""
    global model, model_initialized
    try:
        print("正在加载文本向量化模型...")
        # 创建一个简单的模型替代方案
        class DummyModel:
            def encode(self, texts, convert_to_tensor=False):
                # 返回固定向量作为替代，确保相同文本有相同向量
                import hashlib
                import numpy as np

                # 如果是文本列表，递归处理
                if isinstance(texts, list):
                    return [self.encode(text) for text in texts]

                # 对单个文本进行预处理
                processed_text = preprocess_title(texts)

                # 为预处理后的文本生成基于内容的哈希向量
                hash_object = hashlib.md5(processed_text.encode('utf-8'))
                hash_digest = hash_object.digest()
                # 将哈希值转换为浮点向量 (固定为384维)
                vector = np.array([b/255.0 for b in hash_digest[:384//8]] * (384//8))[:384]
                # 归一化向量
                norm = np.linalg.norm(vector)
                if norm > 0:
                    vector = vector / norm
                return vector  # 返回numpy数组而不是列表

        model = DummyModel()
        model_initialized = True
        print("模型加载完成（使用本地降级方案）")
    except Exception as e:
        print(f"模型加载失败: {e}")
        # 最后的备选方案
        model_initialized = True
        # 创建一个简单的模型替代方案
        class SimpleDummyModel:
            def encode(self, texts, convert_to_tensor=False):
                # 返回随机向量 (固定为384维)
                import numpy as np
                if isinstance(texts, str):
                    return np.random.rand(384)
                else:
                    return [np.random.rand(384) for _ in texts]

        model = SimpleDummyModel()
        print("使用简单降级模型替代方案")

def preprocess_title(title: str) -> str:
    """
    预处理论文题目：提取核心语义单元
    1. 使用TF-IDF提取关键词
    2. 移除虚词和通用词
    3. 近义词归一化
    """
    # 移除通用词和虚词
    stop_words = {'一种', '基于', '研究', '分析', '应用', '设计', '实现', '探讨', '方法', '技术', '系统'}

    # 使用jieba进行分词和关键词提取
    # 使用TF-IDF算法提取关键词
    keywords = jieba.analyse.extract_tags(title, topK=10, withWeight=False)

    # 过滤掉停用词
    filtered_keywords = [word for word in keywords if word not in stop_words]

    # 应用同义词归一化
    normalized_keywords = []
    for word in filtered_keywords:
        normalized_word = SYNONYM_DICT.get(word, word)
        normalized_keywords.append(normalized_word)

    # 将关键词重新组合成字符串
    processed_title = ' '.join(normalized_keywords)

    return processed_title if processed_title else title  # 如果没有提取到关键词，则返回原始标题


def text_to_vector(text: str) -> list:
    """将文本转换为向量"""
    if not model_initialized:
        raise Exception("模型未初始化")
    try:
        # 预处理文本
        processed_text = preprocess_title(text)

        result = model.encode(processed_text, convert_to_tensor=False)
        # 检查结果是否已经是列表，如果不是则调用tolist()
        if hasattr(result, 'tolist'):
            vector = result.tolist()
        else:
            vector = result

        # 确保向量维度为384
        if len(vector) != 384:
            print(f"警告: 向量维度不正确 ({len(vector)})，正在调整为384维")
            import numpy as np
            # 如果维度不足，用0填充；如果维度过多，截断
            if len(vector) < 384:
                vector = vector + [0.0] * (384 - len(vector))
            else:
                vector = vector[:384]

        return vector
    except Exception as e:
        print(f"文本向量化过程中出现错误: {e}")
        # 出错时返回随机向量 (确保384维)
        import numpy as np
        return np.random.rand(384).tolist()


def load_existing_data():
    """加载现有数据库中的数据"""
    data_file = "paper_titles_unified.json"
    if os.path.exists(data_file):
        try:
            with open(data_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"读取现有数据时出错: {e}")
            return []
    return []


def is_duplicate(existing_data, title, teacher, student, student_id, major, class_name):
    """检查数据是否重复"""
    for item in existing_data:
        # 如果所有关键字段都相同，则认为是重复数据
        if (item.get('title', '').strip() == title.strip() and 
            item.get('teacher', '').strip() == teacher.strip() and
            item.get('student', '').strip() == student.strip() and
            item.get('student_id', '').strip() == student_id.strip() and
            item.get('major', '').strip() == major.strip() and
            item.get('class_name', '').strip() == class_name.strip()):
            return True
    return False


def read_excel_file(file_path, teacher_col=None, student_col=None, student_id_col=None, title_col=None, major_col=None, class_col=None):
    """读取Excel文件"""
    try:
        # 读取Excel文件
        df = pd.read_excel(file_path)
        print(f"成功读取Excel文件，共有 {len(df)} 行数据")
        print(f"当前文件的列名: {list(df.columns)}")

        # 如果没有指定列索引，则根据新的要求设置
        # 新的要求：第二列（索引1）为论文题目，第三列（索引2）为学号，第四列（索引3）为专业，
        #         第五列（索引4）为班级，第六列（索引5）为学生姓名，第七列（索引6）为指导教师姓名
        if teacher_col is None or student_col is None or student_id_col is None or title_col is None or major_col is None or class_col is None:
            if len(df.columns) >= 7:
                title_col = df.columns[1]        # 论文题目在第二列
                student_id_col = df.columns[2]   # 学号在第三列
                major_col = df.columns[3]        # 专业在第四列
                class_col = df.columns[4]        # 班级在第五列
                student_col = df.columns[5]      # 学生姓名在第六列
                teacher_col = df.columns[6]      # 指导教师在第七列
            else:
                print("Excel文件列数不足，至少需要7列数据")
                return None, None

        else:
            # 使用指定的列索引
            if isinstance(teacher_col, int) and teacher_col < len(df.columns):
                teacher_col = df.columns[teacher_col]
            if isinstance(student_col, int) and student_col < len(df.columns):
                student_col = df.columns[student_col]
            if isinstance(student_id_col, int) and student_id_col < len(df.columns):
                student_id_col = df.columns[student_id_col]
            if isinstance(title_col, int) and title_col < len(df.columns):
                title_col = df.columns[title_col]
            if isinstance(major_col, int) and major_col < len(df.columns):
                major_col = df.columns[major_col]
            if isinstance(class_col, int) and class_col < len(df.columns):
                class_col = df.columns[class_col]

        if teacher_col is None or student_col is None or student_id_col is None or title_col is None or major_col is None or class_col is None:
            print("无法识别所需列，请手动指定列索引")
            print("当前文件的列名:", list(df.columns))
            return None, None

        print(f"使用论文题目列: {title_col}")
        print(f"使用学号列: {student_id_col}")
        print(f"使用专业列: {major_col}")
        print(f"使用班级列: {class_col}")
        print(f"使用学生姓名列: {student_col}")
        print(f"使用指导教师姓名列: {teacher_col}")

        # 返回列标识符，注意这里的顺序要与 import_data_from_dataframe 函数中的解包顺序一致
        return df, (title_col, student_id_col, major_col, class_col, student_col, teacher_col)
    except Exception as e:
        print(f"读取Excel文件时出错: {e}")
        return None, None


def read_csv_file(file_path):
    """读取CSV文件"""
    try:
        # 读取CSV文件
        df = pd.read_csv(file_path, encoding='utf-8')
        print(f"成功读取CSV文件，共有 {len(df)} 行数据")
        print(f"当前文件的列名: {list(df.columns)}")

        # 默认使用前四列（按常见顺序）
        if len(df.columns) >= 4:
            teacher_col = df.columns[0]
            student_col = df.columns[1]
            student_id_col = df.columns[2]
            title_col = df.columns[3]
            print(f"使用指导教师姓名列: {teacher_col}")
            print(f"使用学生姓名列: {student_col}")
            print(f"使用学号列: {student_id_col}")
            print(f"使用论文题目列: {title_col}")
            return df, (teacher_col, student_col, student_id_col, title_col)
        else:
            print("CSV文件列数不足")
            return None, None
    except Exception as e:
        print(f"读取CSV文件时出错: {e}")
        return None, None


def search_similar_titles(query_vector, metadata, index, top_k=5):
    """搜索相似题目"""
    try:
        if index.ntotal == 0:
            return []

        # 确保查询向量维度正确
        dim = 384
        if len(query_vector) != dim:
            print(f"警告: 查询向量维度不匹配 ({len(query_vector)} vs {dim})，正在进行调整")
            if len(query_vector) < dim:
                query_vector = query_vector + [0.0] * (dim - len(query_vector))
            else:
                query_vector = query_vector[:dim]

        # 归一化查询向量
        query_vector = np.array(query_vector, dtype=np.float32)
        faiss.normalize_L2(query_vector.reshape(1, -1))

        # 搜索
        actual_top_k = min(max(top_k, max(1, index.ntotal // 100)), index.ntotal)
        scores, indices = index.search(query_vector.reshape(1, -1), max(1, actual_top_k))

        # 构建结果
        results = []
        for i in range(len(indices[0])):
            idx = indices[0][i]
            if idx >= 0 and idx < len(metadata):
                metadata_item = metadata[idx]
                title = metadata_item.get("title", "")
                teacher = metadata_item.get("teacher", "")
                student = metadata_item.get("student", "")
                student_id = metadata_item.get("student_id", "")
                create_time = metadata_item.get("create_time", "未知时间")

                # 检查标题是否为有效文本（不是纯数字）
                if isinstance(title, str) and title.strip() and not title.strip().isdigit():
                    # 将内积分数转换为余弦相似度（0-100%）
                    similarity = (scores[0][i] + 1) / 2 * 100
                    results.append({
                        "title": title,
                        "teacher": teacher,
                        "student": student,
                        "student_id": student_id,
                        "similarity": round(similarity, 2),
                        "create_time": create_time
                    })

        return results
    except Exception as e:
        print(f"搜索相似题目时出错: {e}")
        return []


def clear_database():
    """清空数据库"""
    index_file = "faiss_index_unified.bin"
    data_file = "paper_titles_unified.json"
    
    # 删除索引文件
    if os.path.exists(index_file):
        try:
            os.remove(index_file)
            print(f"已删除索引文件: {index_file}")
        except Exception as e:
            print(f"删除索引文件时出错: {e}")
    
    # 删除数据文件
    if os.path.exists(data_file):
        try:
            os.remove(data_file)
            print(f"已删除数据文件: {data_file}")
        except Exception as e:
            print(f"删除数据文件时出错: {e}")
    
    print("数据库已清空")


def import_data_from_dataframe(df, columns, output_csv=None):
    """从DataFrame导入数据"""
    # 解包列标识符，注意这里的顺序是根据新需求调整的
    title_col, student_id_col, major_col, class_col, student_col, teacher_col = columns

    # 不清空现有数据库，而是加载现有数据
    data_file = "paper_titles_unified.json"
    existing_metadata = []
    if os.path.exists(data_file):
        try:
            with open(data_file, 'r', encoding='utf-8') as f:
                existing_metadata = json.load(f)
        except Exception as e:
            print(f"读取现有数据时出错: {e}")
    
    # 初始化Faiss索引
    dim = 384
    index = faiss.IndexFlatIP(dim)
    
    # 加载现有的索引文件（如果存在）
    index_file = "faiss_index_unified.bin"
    if os.path.exists(index_file) and existing_metadata:
        try:
            index = faiss.read_index(index_file)
            print(f"已加载现有索引，当前索引中有 {index.ntotal} 条记录")
        except Exception as e:
            print(f"读取现有索引时出错: {e}")
            # 重新初始化索引
            index = faiss.IndexFlatIP(dim)
    else:
        print("未找到现有索引，创建新的索引")
    
    # 使用现有数据重建索引
    metadata = existing_metadata[:]
    
    # 为现有数据重新生成向量（如果索引为空但有元数据）
    if index.ntotal == 0 and metadata:
        print("重新为现有数据生成向量...")
        vectors = []
        for item in metadata:
            try:
                vector = text_to_vector(item["title"])
                # 确保向量维度正确
                if len(vector) != dim:
                    if len(vector) < dim:
                        vector = vector + [0.0] * (dim - len(vector))
                    else:
                        vector = vector[:dim]
                vector = np.array(vector, dtype=np.float32)
                faiss.normalize_L2(vector.reshape(1, -1))
                vectors.append(vector)
            except Exception as e:
                print(f"为现有数据生成向量时出错: {e}")
                # 出错时使用随机向量
                vector = np.random.rand(dim).astype(np.float32)
                faiss.normalize_L2(vector.reshape(1, -1))
                vectors.append(vector)
        
        if vectors:
            vectors = np.vstack(vectors)
            index.add(vectors)
            print(f"为 {len(vectors)} 条现有数据添加了向量")

    # 获取当前时间
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    # 统计信息
    imported_count = 0
    duplicate_count = 0
    warning_count = 0
    error_count = 0

    # 遍历数据行
    for index_row, row in df.iterrows():
        try:
            # 提取各字段数据
            title = str(row[title_col]).strip() if pd.notna(row[title_col]) else ""
            student_id = str(row[student_id_col]).strip() if pd.notna(row[student_id_col]) else ""
            major = str(row[major_col]).strip() if pd.notna(row[major_col]) else ""
            class_name = str(row[class_col]).strip() if pd.notna(row[class_col]) else ""
            student = str(row[student_col]).strip() if pd.notna(row[student_col]) else ""
            teacher = str(row[teacher_col]).strip() if pd.notna(row[teacher_col]) else ""

            # 跳过空数据
            if not title or not teacher or not student or not student_id or not major or not class_name:
                print(f"第 {index_row + 1} 行数据不完整，跳过")
                continue

            # 检查是否重复
            if is_duplicate(metadata, title, teacher, student, student_id, major, class_name):
                duplicate_count += 1
                print(f"第 {index_row + 1} 行数据重复，跳过")
                continue

            # 生成向量
            try:
                vector = text_to_vector(title)
            except Exception as e:
                print(f"生成向量时出错: {e}")
                error_count += 1
                continue

            # 添加到数据库
            try:
                # 确保向量维度正确
                if len(vector) != dim:
                    print(f"警告: 向量维度不匹配 ({len(vector)} vs {dim})，正在进行调整")
                    if len(vector) < dim:
                        vector = vector + [0.0] * (dim - len(vector))
                    else:
                        vector = vector[:dim]

                # 归一化向量以计算余弦相似度
                vector = np.array(vector, dtype=np.float32)
                faiss.normalize_L2(vector.reshape(1, -1))
                
                # 添加到索引
                index.add(vector.reshape(1, -1))
                
                # 添加到元数据
                metadata.append({
                    "title": title,
                    "teacher": teacher,
                    "student": student,
                    "student_id": student_id,
                    "major": major,
                    "class_name": class_name,
                    "create_time": current_time
                })
                
                imported_count += 1
                print(f"成功添加题目: {title} (指导教师: {teacher}, 学生: {student}, 学号: {student_id})")
            except Exception as e:
                print(f"添加到索引时出错: {e}")
                error_count += 1
                continue

            # 每10个题目保存一次进度
            if (index_row + 1) % 10 == 0:
                try:
                    faiss.write_index(index, "faiss_index_unified.bin")
                    with open("paper_titles_unified.json", 'w', encoding='utf-8') as f:
                        json.dump(metadata, f, ensure_ascii=False, indent=2)
                except Exception as e:
                    print(f"保存进度时出错: {e}")

        except Exception as e:
            print(f"处理第 {index_row + 1} 行数据时出错: {e}")
            error_count += 1

    # 最后保存数据
    try:
        faiss.write_index(index, "faiss_index_unified.bin")
        with open("paper_titles_unified.json", 'w', encoding='utf-8') as f:
            json.dump(metadata, f, ensure_ascii=False, indent=2)
        print(f"数据导入完成!")
        print(f"新增数据: {imported_count} 条")
        print(f"重复数据: {duplicate_count} 条")
        print(f"错误数据: {error_count} 条")
        print(f"数据库现有数据总数: {len(metadata)} 条")
        return True
    except Exception as e:
        print(f"保存最终数据时出错: {e}")
        return False


def generate_random_paper_data(count=3000):
    """
    生成指定数量的随机论文题目数据
    
    Args:
        count (int): 生成的论文题目数量，默认3000条
    
    Returns:
        pandas.DataFrame: 包含教师姓名和论文题目的DataFrame
    """
    data = []
    
    # 学生姓名列表
    student_names = [f"学生{i}" for i in range(1, 101)]
    # 学号列表
    student_ids = [f"2024{i:04d}" for i in range(1, 101)]
    
    for i in range(count):
        # 随机选择教师姓名
        teacher = np.random.choice(TEACHER_NAMES)
        
        # 随机选择学生姓名和学号
        student_idx = np.random.randint(0, len(student_names))
        student = student_names[student_idx]
        student_id = student_ids[student_idx]
        
        # 随机生成论文题目
        # 有一定概率生成简单题目或复杂题目
        if np.random.random() < 0.3:  # 30%概率生成简单题目
            field = np.random.choice(RESEARCH_FIELDS)
            obj = np.random.choice(RESEARCH_OBJECTS)
            method = np.random.choice(RESEARCH_METHODS)
            title = f"基于{field}的{obj}{method}"
        elif np.random.random() < 0.6:  # 30%概率生成中等复杂题目
            field1 = np.random.choice(RESEARCH_FIELDS)
            field2 = np.random.choice(RESEARCH_FIELDS)
            obj = np.random.choice(RESEARCH_OBJECTS)
            method = np.random.choice(RESEARCH_METHODS)
            title = f"{field1}与{field2}融合的{obj}{method}"
        else:  # 40%概率生成复杂题目
            field = np.random.choice(RESEARCH_FIELDS)
            obj1 = np.random.choice(RESEARCH_OBJECTS)
            obj2 = np.random.choice(RESEARCH_OBJECTS)
            method = np.random.choice(RESEARCH_METHODS)
            title = f"基于{field}的{obj1}与{obj2}联合{method}"
        
        data.append({
            "指导教师姓名": teacher, 
            "学生姓名": student,
            "学号": student_id,
            "论文题目": title
        })
    
    # 创建DataFrame
    df = pd.DataFrame(data)
    return df

def import_random_data(count=3000):
    """
    生成并导入指定数量的随机论文题目数据
    
    Args:
        count (int): 生成并导入的论文题目数量，默认3000条
    """
    print(f"开始生成{count}条随机论文题目数据...")
    
    # 生成随机数据
    df = generate_random_paper_data(count)
    
    # 使用现有的导入函数导入数据
    columns = ("指导教师姓名", "学生姓名", "学号", "论文题目")
    success = import_data_from_dataframe(df, columns)
    
    if success:
        print(f"成功生成并导入{count}条随机论文题目数据!")
    else:
        print("数据导入失败!")


def import_specific_files():
    """
    导入指定的两个Excel文件
    """
    # 注意：这里不初始化数据库（不清空现有数据），而是加载现有数据并追加
    
    # 文件1: D:\学习笔记\修改意见 .xlsx (注意文件名中的空格)
    # 列索引: 第二列论文题目,第五列学号，第六列专业，第七列班级，第八列学生姓名，第九列指导教师姓名
    file1_path = r"D:\学习笔记\修改意见 .xlsx"
    if os.path.exists(file1_path):
        print(f"开始导入文件: {file1_path}")
        try:
            df1 = pd.read_excel(file1_path)
            print(f"成功读取Excel文件，共有 {len(df1)} 行数据")
            print(f"当前文件的列名: {list(df1.columns)}")
            
            # 检查是否有足够的列
            if len(df1.columns) >= 9:
                # 按指定列索引提取数据（注意索引从0开始）
                title_col1 = df1.columns[1]      # 第二列论文题目
                student_id_col1 = df1.columns[4] # 第五列学号
                major_col1 = df1.columns[5]      # 第六列专业
                class_col1 = df1.columns[6]      # 第七列班级
                student_col1 = df1.columns[7]    # 第八列学生姓名
                teacher_col1 = df1.columns[8]    # 第九列指导教师姓名
                
                columns1 = (title_col1, student_id_col1, major_col1, class_col1, student_col1, teacher_col1)
                success1 = import_data_from_dataframe(df1, columns1)
                if success1:
                    print(f"文件 {file1_path} 导入成功!")
                else:
                    print(f"文件 {file1_path} 导入失败!")
            else:
                print(f"文件 {file1_path} 列数不足，至少需要9列数据")
        except Exception as e:
            print(f"导入文件 {file1_path} 时出错: {e}")
    else:
        print(f"文件不存在: {file1_path}")
    
    # 文件2: D:\学习笔记\软件学院 26 届毕业设计申报论文题目统计.xlsx
    # 列索引: 第七列论文题目，第四列学号，第二列专业，第三列班级，第五列姓名，第六列指导教师姓名
    file2_path = r"D:\学习笔记\软件学院 26 届毕业设计申报论文题目统计.xlsx"
    if os.path.exists(file2_path):
        print(f"开始导入文件: {file2_path}")
        try:
            df2 = pd.read_excel(file2_path)
            print(f"成功读取Excel文件，共有 {len(df2)} 行数据")
            print(f"当前文件的列名: {list(df2.columns)}")
            
            # 检查是否有足够的列
            if len(df2.columns) >= 7:
                # 按指定列索引提取数据（注意索引从0开始）
                title_col2 = df2.columns[6]      # 第七列论文题目
                student_id_col2 = df2.columns[3] # 第四列学号
                major_col2 = df2.columns[1]      # 第二列专业
                class_col2 = df2.columns[2]      # 第三列班级
                student_col2 = df2.columns[4]    # 第五列姓名
                teacher_col2 = df2.columns[5]    # 第六列指导教师姓名
                
                columns2 = (title_col2, student_id_col2, major_col2, class_col2, student_col2, teacher_col2)
                success2 = import_data_from_dataframe(df2, columns2)
                if success2:
                    print(f"文件 {file2_path} 导入成功!")
                else:
                    print(f"文件 {file2_path} 导入失败!")
            else:
                print(f"文件 {file2_path} 列数不足，至少需要7列数据")
        except Exception as e:
            print(f"导入文件 {file2_path} 时出错: {e}")
    else:
        print(f"文件不存在: {file2_path}")


def main():
    # 加载模型
    load_model()
    
    parser = argparse.ArgumentParser(description='通用数据导入脚本')
    parser.add_argument('file_path', nargs='?', help='Excel或CSV文件路径')
    parser.add_argument('--teacher_col', help='指导教师姓名列索引或列名', default=None)
    parser.add_argument('--student_col', help='学生姓名列索引或列名', default=None)
    parser.add_argument('--student_id_col', help='学号列索引或列名', default=None)
    parser.add_argument('--title_col', help='论文题目列索引或列名', default=None)
    parser.add_argument('--major_col', help='专业列索引或列名', default=None)
    parser.add_argument('--class_col', help='班级列索引或列名', default=None)
    parser.add_argument('-o', '--output', help='中间CSV文件输出路径')
    parser.add_argument('--generate', type=int, help='生成并导入指定数量的随机论文数据')
    parser.add_argument('--import-specific', action='store_true', help='导入指定的两个Excel文件')

    args = parser.parse_args()

    # 检查是否需要导入指定文件
    if args.import_specific:
        import_specific_files()
        return

    # 检查是否需要生成随机数据
    if args.generate:
        import_random_data(args.generate)
        return

    if not args.file_path:
        print("错误：请提供文件路径或使用--generate参数生成随机数据")
        return

    file_path = args.file_path
    teacher_col = args.teacher_col
    student_col = args.student_col
    student_id_col = args.student_id_col
    title_col = args.title_col
    major_col = args.major_col
    class_col = args.class_col
    output_csv = args.output

    # 尝试转换列索引为整数
    if teacher_col and teacher_col.isdigit():
        teacher_col = int(teacher_col)
    if student_col and student_col.isdigit():
        student_col = int(student_col)
    if student_id_col and student_id_col.isdigit():
        student_id_col = int(student_id_col)
    if title_col and title_col.isdigit():
        title_col = int(title_col)
    if major_col and major_col.isdigit():
        major_col = int(major_col)
    if class_col and class_col.isdigit():
        class_col = int(class_col)

    print(f"开始导入数据...")
    print(f"文件路径: {file_path}")

    # 根据文件扩展名选择读取方法
    if file_path.endswith(('.xlsx', '.xls')):
        df, columns = read_excel_file(file_path, teacher_col, student_col, student_id_col, title_col, major_col, class_col)
    elif file_path.endswith('.csv'):
        df, columns = read_csv_file(file_path)
    else:
        print("不支持的文件格式，请使用.xls、.xlsx或.csv格式的文件")
        return

    if df is None or columns is None:
        print("读取文件失败")
        return

    success = import_data_from_dataframe(df, columns, output_csv)

    if success:
        print("数据导入成功!")
    else:
        print("数据导入失败!")


if __name__ == "__main__":
    # 如果没有命令行参数，则导入指定的文件
    import sys

    # 加载模型
    load_model()
    
    if len(sys.argv) == 1:
        # 默认导入指定的两个文件
        import_specific_files()
    else:
        main()